Skip to content

[Nemotron][CUDA] Add deterministic moe_router_dispatch - #504

Draft
RichApple123 wants to merge 1 commit into
RL-Align:test-nemotronfrom
RichApple123:feat/nemotron-router-sm90
Draft

RichApple123 wants to merge 1 commit into
RL-Align:test-nemotronfrom
RichApple123:feat/nemotron-router-sm90

Conversation

@RichApple123

Copy link
Copy Markdown

Implements the SM90 CUDA moe_router_dispatch work item of #434: FP32 projection, sigmoid, bias-corrected top-six selection, normalization and deterministic token packing, with backward for routing weights and packed payloads. Supports BF16/FP32 inputs with FP32 router weights at Nano's 2688-hidden, 128-expert geometry.

Adds a numerical reference, CPU/GPU correctness and batch-invariance tests, a cuBLAS + Transformer Engine benchmark, documentation and CI wiring for test-nemotron. The ordering and fixed arithmetic policy are proposed for section 4.4 review. Expert MLP/combine and distributed communication are outside this change.

Validation:

  • Single-H20 qualification of the unchanged operator: 106 tests passed; output and gradient bits matched the prior schedule in 36 cases; four Compute Sanitizer tools passed.
  • After consolidating the GPU test files, local checks: 25 passed, 81 CUDA skipped. Original test function bodies and parameterization are preserved; no new GPU run was performed for this consolidation. Formatting and lint checks pass.
  • Random BF16 forward+backward at 8192/32768 tokens is 1.08-1.09x faster than cuBLAS + TE 2.19 index dispatch/autograd. FP32 is near parity; small-token regressions are reported. Both providers use eager timing. The benchmark records raw timing samples, incremental peak memory and source hashes; the measured matrix is retained separately for review.

The target branch still has the eight documentation link warnings addressed by #458 on main. This change adds no documentation warnings; strict docs passes locally with a separate fix for those links. Target-branch CI and maintainer approval of the proposed contract remain pending.

Signed-off-by: RichApple123 <134492660+RichApple123@users.noreply.github.com>
@coderabbitai

coderabbitai Bot commented Oct 10, 2026

Copy link
Copy Markdown

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

  • Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true
  • Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Comment @coderabbitai help to get the list of available commands.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant